import itertools

import torch


def all_combinations(lst):
    # 存储所有组合
    all_combs = []
    # 遍历所有可能的组合长度
    for r in range(1, len(lst) + 1):
        # 获取当前长度的所有组合，并添加到总列表中
        combs = list(itertools.combinations(lst, r))
        all_combs.extend(combs)
    return all_combs


def randomSelectFromTensor(index_tensor_list, select_num):
    """
    从给定的张量中随机选择指定数量的元素。

    :param index_tensor_list: 包含索引的张量
    :param select_num: 需要选择的元素数量
    :return: 包含随机选择的索引的张量
    """
    # 确保输入是一个张量
    if not isinstance(index_tensor_list, torch.Tensor):
        raise ValueError("index_tensor_list must be a torch.Tensor")

    # 获取张量的长度
    tensor_length = index_tensor_list.size(0)

    # 如果需要选择的数量大于张量的长度，则返回整个张量
    if select_num >= tensor_length:
        return index_tensor_list

    # 生成一个随机排列的索引
    random_indices = torch.randperm(tensor_length)

    # 选择前select_num个索引
    selected_indices = index_tensor_list[random_indices[:select_num]]

    return selected_indices



def weightedRandomSelect(weights, num_samples, replacement=False):
    """
    根据给定的权重从张量中进行抽样。

    :param weights: 一个一维张量，表示每个元素被选中的概率权重
    :param num_samples: 要选择的样本数量
    :param replacement: 是否放回抽样，默认为False（不放回）
    :return: 一个一维张量，包含抽样的索引
    """
    if not isinstance(weights, torch.Tensor):
        raise ValueError("weights must be a torch.Tensor")

    # 确保weights是一个一维张量
    if weights.dim() != 1:
        raise ValueError("weights must be a one-dimensional tensor")

    # 使用torch.multinomial进行加权抽样
    sampled_indices = torch.multinomial(weights, num_samples, replacement=replacement)

    return sampled_indices

